# encoding: utf-8

import json
import numpy as np
from typing import List
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

from zcbot_web_core.model.handler import Resp

from modelscope.pipelines import pipeline

pretrained_model = 'damo/nlp_raner_named-entity-recognition_chinese-base-ecom-50cls'
ner_pipeline = pipeline('named-entity-recognition', pretrained_model)

app = FastAPI()


# Extend the JSONEncoder class
class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)


class Item(BaseModel):
    text: str = None
    textList: List[str] = None


def ner_one(text: str):
    result = ner_pipeline(text, )
    result = json.loads(json.dumps(result, cls=NpEncoder))
    return result


def ner_batch(text_list: List[str]):
    result = ner_pipeline(text_list)
    result = json.loads(json.dumps(result, cls=NpEncoder))
    return result


@app.get("/ali-ner")
def get_ner(text: str):
    result = ner_one(text)
    return Resp.ok(result)


@app.post("/ali-ner")
def post_ner(item: Item):
    if item.text:
        result = ner_one(item.text)
    elif item.textList:
        result = ner_batch(item.textList)
    else:
        return Resp.err("body cannot be null")
    return Resp.ok(result)


if __name__ == '__main__':
    uvicorn.run(app, host="0.0.0.0", port=6698)
